# 模型组网的功能是将层串起来，实现数据的前向传播和梯度的反向传播
# 添加层的时候，按照顺序添加层的参数
# Sequential方法说明：
#   add: 向组网中添加层
#   forward: 按照组网构建的层顺序，依次前向传播
#   backward: 接收损失函数的梯度，按照层的逆序反向传播
from 神经网络原理.fc_test import Layer

'''
组网实体类 对线性层或者卷积层进行遍历调用layer的前向传播或者后向传播的grad
'''
class Sequential:
    def __init__(self, *args, **kwargs):
        self.graphs = []
        self._parameters = []
        for arg_layer in args:
            if isinstance(arg_layer, Layer):
                self.graphs.append(arg_layer)
                self._parameters += arg_layer.parameters()

    def add(self, layer):
        assert isinstance(layer, Layer), "The type of added layer must be Layer, but got {}.".format(type(layer))
        self.graphs.append(layer)
        self._parameters += layer.parameters()

    def forward(self, x):
        for graph in self.graphs:
            x = graph(x)
        return x

    def backward(self, grad):
        # grad backward in inverse order of graph
        for graph in self.graphs[::-1]:
            grad = graph.backward(grad)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def __str__(self):
        string = 'Sequential:\n'
        for graph in self.graphs:
            string += graph.__str__() + '\n'
        return string

    def parameters(self):
        return self._parameters